#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <arpa/inet.h>

#include "dns_server_rr_files.h"
#include "dns_packet_parsing.h"
#include "dns_util.h"
#include "local_dns_util.h"

/*
 * Global variables
 */

SOCKET server_socket = INVALID_SOCKET;

/*
 * Functions
 */

/* Ask another server to answer the question */
void ask_remote(char * buffer, int msg_length, SOCKADDR_IN * remote_addr, socklen_t remote_addr_len, char * remote_addr_str, int remote_port) {
	remote_addr->sin_port = htons(remote_port);

	if (!inet_pton(AF_INET, remote_addr_str, &(remote_addr->sin_addr)))
		exit_runerr("Cannot convert address", EXIT_FAILURE, &server_socket, NULL);

	if (sendto(server_socket, buffer, msg_length, 0, (SOCKADDR *) remote_addr, remote_addr_len) < 0)
		exit_syserr("sendto()", &server_socket, NULL);
}

/* Wait to receive the answer of a server */
void wait_remote(char * buffer, int * msg_length, SOCKADDR_IN * remote_addr, socklen_t * remote_addr_len) {
	if ((*msg_length = recvfrom(server_socket, buffer, BUF_SIZE, 0, (SOCKADDR *) remote_addr, remote_addr_len)) < 0)
		exit_syserr("recvfrom()", &server_socket, NULL);
}

/*
 * Main
 */
int main(int argc, char ** argv) {

	int local_port;

	/* Read arguments */
	if (argc != 2) {
		fprintf(stderr, "Usage: %s local_port\n", argv[0]);
		exit(EXIT_FAILURE);
	}

	local_port = atoi(argv[1]);

	/* --- */

	SOCKADDR_IN local_addr;
	socklen_t local_addr_len;
	SOCKADDR_IN client_addr;
	socklen_t client_addr_len;
	SOCKADDR_IN remote_addr;
	socklen_t remote_addr_len;

	/* Create socket  */
	if ((server_socket = socket(AF_INET, SOCK_DGRAM, 0)) == INVALID_SOCKET)
		exit_syserr("socket()", NULL);

  /* Init local server address */
	local_addr_len = sizeof(local_addr);
	memset(&local_addr, 0, local_addr_len);
	local_addr.sin_family = AF_INET;
	local_addr.sin_port = htons(local_port);
	local_addr.sin_addr.s_addr = htonl(INADDR_ANY);

	/* Init client address */
	client_addr_len = sizeof(client_addr);
	memset(&client_addr, 0, client_addr_len);

	/* Init remote server address */
	remote_addr_len = sizeof(remote_addr);
	memset(&remote_addr, 0, remote_addr_len);
	remote_addr.sin_family = AF_INET;

	/* Bind socket */
	if (bind(server_socket, (SOCKADDR *) &local_addr, local_addr_len) == SOCKET_ERROR)
		exit_syserr("bind()", &server_socket, NULL);

	/* Run server loop */
	char buffer[BUF_SIZE];
	int msg_length;
	char remote_addr_str[IPV4_STR_LENGTH];
	int remote_port;
	while (true) {
		/* Reset buffer */
		memset(buffer, 0, BUF_SIZE);

		/* Wait for a client request */
		printf("\nWaiting...\n");

		/* Read the message */
		if ((msg_length = recvfrom(server_socket, buffer, BUF_SIZE, 0, (SOCKADDR *) &client_addr, &client_addr_len)) < 0)
			exit_syserr("recvfrom()", &server_socket, NULL);
		printf("Message received\n");

		/* Extract requested name */
		char qname[DOMAIN_STR_LENGTH];
		if (qname_from_question(buffer, BUF_SIZE, qname) < 0) {
			fprintf(stderr, "Cannot read qname\n");
			continue;
		}
		printf("QName: %s\n", qname);

		/* Ask the ROOT server */
		printf("Asking ROOT server\n");
		ask_remote(buffer, msg_length, &remote_addr, remote_addr_len, ROOT_ADDR, ROOT_PORT);

		/* Wait for ROOT server to answer */
		printf("Waiting for ROOT server answer\n");
		wait_remote(buffer, &msg_length, &remote_addr, &remote_addr_len);

		/* Get the TLD server's address */
		extract_remote_info(buffer, remote_addr_str, &remote_port);
		printf("TLD server is at %s:%d\n", remote_addr_str, remote_port);

		/* Ask the TLD server */
		printf("Asking TLD server\n");
		ask_remote(buffer, msg_length, &remote_addr, remote_addr_len, remote_addr_str, remote_port);

		/* Wait for TLD server to answer */
		printf("Waiting for TLD server answer\n");
		wait_remote(buffer, &msg_length, &remote_addr, &remote_addr_len);

		/* Get the AUTHORITATIVE server's address */
		extract_remote_info(buffer, remote_addr_str, &remote_port);
		printf("AUTHORITATIVE server is at %s:%d\n", remote_addr_str, remote_port);

		/* Ask the AUTHORITATIVE server */
		printf("Asking AUTHORITATIVE server\n");
		ask_remote(buffer, msg_length, &remote_addr, remote_addr_len, remote_addr_str, remote_port);

		/* Wait for AUTHORITATIVE server to answer */
		printf("Waiting for AUTHORITATIVE server answer\n");
		wait_remote(buffer, &msg_length, &remote_addr, &remote_addr_len);

		/* Send answer to back to the client */
		printf("Answering to the client\n");
		if (sendto(server_socket, buffer, msg_length, 0, (SOCKADDR *) &client_addr, client_addr_len) < 0)
			exit_syserr("sendto()", &server_socket, NULL);
	}

	/* --- */

	close_socket(&server_socket);

	return EXIT_SUCCESS;
}
